package skewreduce.framework;

import java.io.IOException;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.util.ReflectionUtils;

public class MultipleOutputs {
    private static final String MULTIOUTPUT_NAME = "multipleoutput.names";
    private static final String MULTIOUTPUT_CLASS= "multipleoutput.classes";

    private Map<String,RecordWriter<?,?>> writers;
    private List<RecordWriter<?,?>> writersList;

    public static void addOutput(Job job,
            String name,Class<? extends MultiOutputFormat<?,?>> writerClass)
    throws IllegalArgumentException {
        addOutput(job.getConfiguration(),name,writerClass);
    }


    public static void addOutput(Configuration conf,
            String name,Class<? extends MultiOutputFormat<?,?>> writerClass)
    throws IllegalArgumentException {
        // check configuration
        String names = conf.get(MULTIOUTPUT_NAME); 
        if ( names == null ) {
            conf.set(MULTIOUTPUT_NAME,name);
            conf.set(MULTIOUTPUT_CLASS,writerClass.getName());
        } else {
            conf.set(MULTIOUTPUT_NAME,names+","+name);
            String classes = conf.get(MULTIOUTPUT_CLASS);
            conf.set(MULTIOUTPUT_CLASS,classes+","+writerClass.getName());
        }
    }

    public MultipleOutputs(TaskAttemptContext context)
    throws IOException, InterruptedException {
        Configuration conf = context.getConfiguration();

        String[] names = conf.getStrings(MULTIOUTPUT_NAME);
        Class<?>[] classes = conf.getClasses(MULTIOUTPUT_CLASS);

        if ( names.length != classes.length ) {
            throw new IOException("the number of outputs and classes do not match");
        }

        HashMap<String,RecordWriter<?,?>> tmpOut = new HashMap<String,RecordWriter<?,?>>();

        List<RecordWriter<?,?>> tmpOutList = new ArrayList<RecordWriter<?,?>>(names.length);

        for ( int i = 0; i < names.length; ++i ) {
            MultiOutputFormat<?,?> outFmt = (MultiOutputFormat<?,?>)ReflectionUtils.newInstance(classes[i],conf);
            RecordWriter<?,?> writer = outFmt.getRecordWriter(context,names[i]);
            tmpOut.put( names[i], writer );
            tmpOutList.add(writer);
        }

        writers = Collections.unmodifiableMap(tmpOut);
        writersList = Collections.unmodifiableList(tmpOutList);
    }

    @SuppressWarnings("unchecked")
    public <K,V> void write(String name,K key,V val)
    throws IOException, InterruptedException {
        ((RecordWriter<K,V>)(writers.get(name))).write(key,val);
    }

    @SuppressWarnings("unchecked")
    public <K,V> void write(int idx,K key,V val)
    throws IOException, InterruptedException {
        ((RecordWriter<K,V>)(writersList.get(idx))).write(key,val);
    }

    public void close(TaskAttemptContext context)
    throws IOException,InterruptedException {
        for ( RecordWriter<?,?> w : writersList ) {
            w.close(context);
        }
    }
}
